from functools import partial

import numpy as np
import torch

import sko.operators.mutation


def gpu_atm_costfunc(y, D, aver_1, max_cap, w_0, cutoff_list, I_0, avg_D2, C, std_D2, cutoff_limit,
                     deposit_history=None, withdraw_history=None, args=None):
    if deposit_history is not None and withdraw_history is not None:
        D = withdraw_history - deposit_history
        est_D2 = torch.stack([deposit_history])
    else:
        est_D2 = avg_D2 + C * std_D2

    y0 = y.reshape([y.shape[0], -1, D.shape[1]])
    I, L, cutoff, w = torch.zeros([4, y.shape[0], 7, D.shape[1]], dtype=torch.float32, device=y.device)
    w[:, 0] = w_0.T
    w[:, 1:] = y0 * (max_cap * 10000 / (1 << args.bit_precision))
    w[:, torch.tensor(args.restday).long() - args.day] = 0

    cutoff[:, -1] = cutoff_list.T

    # 初始化每日库存矩阵I，每日缺钞矩阵L_before（换钞前的缺钞值），L_after（换钞后的缺钞值）
    I[:, -1] = I_0.T
    for t in range(7):
        cutoff[:, t] = cutoff[:, t - 1] + 1

        if ((args.day - 1 + t) % 7 + 1) in args.restday:
            I[:, t] = I[:, t - 1] - D[t]

        else:
            I[:, t] = I[:, t - 1] - D[t] * 0.5
            L[:, t] += torch.relu(-I[:, t])
            I[:, t] = torch.relu_(I[:, t])

            I[:, t] = torch.where(w[:, t] == 0, I[:, t], w[:, t])
            cutoff[:, t] *= w[:, t] == 0

            I[:, t] = I[:, t] - D[t] * 0.5

        L[:, t] += torch.relu(-I[:, t])
        I[:, t] = torch.relu_(I[:, t])

    if args.val == True:
        weight = 1
    else:
        weight = torch.tensor([0, 1 / 2, 1 / 3, 1 / 4, 1 / 5, 1 / 6, 1 / 7]).to(w.device).float()
        weight = weight / torch.sum(weight, dim=0) * 7

    over_cutoff = torch.sum(
        torch.relu(cutoff - cutoff_limit.T)[:, 1:],
        (-1, -2)
    )

    I_sum = torch.sum(I, -1)
    I_sum_before = torch.roll(I_sum, shifts=1, dims=1)
    I_sum_before[:, 0] = torch.sum(I_0)
    depos = torch.sum(est_D2 + w, -1) + I_sum_before
    withd = depos - I_sum
    usage_rate = torch.mean(
        withd / depos * weight,
        -1
    )

    short_rate = torch.mean(
        torch.mean(
            L.float() / (D + 1),
            -1
        ) * weight,
        -1
    )

    vaild_day = (7 - (args.day not in args.restday) - len(args.restday)) + 1
    add_freq = torch.count_nonzero(w, (-1, -2)) / (w.shape[-1] * vaild_day)
    return w, I, cutoff, torch.vstack([over_cutoff, usage_rate, short_rate, add_freq])


def parallel_atm_costfunc(y, DATA, args, device):
    myfun = partial(gpu_atm_costfunc,
                    D=torch.from_numpy(DATA.D.astype(np.int32)).int().to(device),
                    aver_1=torch.from_numpy(DATA.aver_1).int().to(device),
                    max_cap=torch.from_numpy(DATA.max_cap.astype(np.int32)).int().to(device),
                    w_0=torch.from_numpy(DATA.w_0).int().to(device),
                    cutoff_list=torch.from_numpy(DATA.cutoff_list).int().to(device),
                    I_0=torch.from_numpy(DATA.I_0).int().to(device),
                    avg_D2=torch.from_numpy(DATA.avg_D2).int().to(device),
                    C=DATA.C,
                    std_D2=torch.from_numpy(DATA.std_D2).int().to(device),
                    cutoff_limit=torch.from_numpy(DATA.cutoff_limit).int().to(device),
                    args=args
                    )

    ans = myfun(y)
    over_cutoff, usage_rate, short_rate, add_freq = ans[3]

    res = 1e6 * (1e0 * torch.relu(over_cutoff - args.limit_cutoff)) + \
          1e0 * (1e3 * torch.relu(add_freq - args.limit_addfreq)).int() + \
          torch.relu(args.limit_usage - usage_rate) + \
          torch.relu(short_rate - args.limit_shortage) * args.ratio_overadd
    return res


# 进化算法
class GA_ATM(sko.GA.GA):

    def __init__(self, func, DATA, args):
        self.args = args
        self.has_constraint = False
        self.max_stay_iter = args.max_stay_step
        self.max_cap = DATA.max_cap

        super().__init__(func, DATA.w_0.size * 6, size_pop=args.batch, max_iter=args.epochs, prob_mut=args.mut_std,
                         early_stop=args.max_stay_step,
                         lb=0, ub=(1 << args.bit_precision) - 1, precision=1)
        self.Chrom |= 1

        cutoff = np.repeat(DATA.cutoff_list.T, args.batch, axis=0) + 1
        cutoff[:, DATA.w_0.reshape(-1) != 0] = 0
        self.Chrom = self.Chrom.reshape(args.batch, 6, DATA.w_0.size)

        for i in range(1, 7):
            cutoff += 1
            if (args.day - 1 + i) % 7 + 1 in args.restday:
                self.Chrom[:, i - 1, :] = 0
            else:
                res = 0
                while (args.day + i + res) % 7 + 1 in args.restday:
                    res += 1

                self.Chrom[:, i - 1, :][cutoff + res <= DATA.cutoff_limit.T] = 0
                cutoff[cutoff + res > DATA.cutoff_limit.T] = 0

        self.Chrom = self.Chrom.reshape(args.batch, -1)
